from torchvision.transforms._presets import ImageClassification
from torchvision import datasets,transforms
from PIL import Image
import torch
import time
from torchsummary import summary
import os
from tqdm import tqdm
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
BATCH_SIZE = 16

PATH_best = 'model/image_classification_best2_adam300.pth'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("设备：", device)

transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),  # 转换为[C, H, W]
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

class ImageClassification(nn.Module):
    def __init__(self):
        super(ImageClassification, self).__init__()
        # 卷积层保持不变
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.pool1 = nn.MaxPool2d(2, 2)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.pool2 = nn.MaxPool2d(2, 2)

        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.pool3 = nn.MaxPool2d(2, 2)

        # conv1后：32x32 → 32x32x32 → 展平维
        # pool1后：32x32 → 16x16x32 → 展平维
        # conv2后：16x16 → 16x16x64 → 展平维
        # pool2后：16x16 → 8x8x64 → 展平维
        # conv3后：8x8 → 8x8x128 → 展平维
        # pool3后：8x8 → 4x4x128 → 展平维

        self.dropout = nn.Dropout(0.2)

        # 全连接层
        self.linear1 = nn.Linear(2048, 512)
        self.linear2 = nn.Linear(512, 128)
        self.linear3 = nn.Linear(128, 32)

        # 输出层
        self.out = nn.Linear(32, 10)

    def forward(self, x):
        # 卷积+relu+池化
        x = self.conv1(x)
        x = self.bn1(x)  # 先BN
        x = torch.relu(x)  # 再激活
        x = self.pool1(x)

        # 卷积+relu+池化
        x = self.conv2(x)
        x = self.bn2(x)  # BN
        x = torch.relu(x)
        x = self.pool2(x)

        # 卷积+relu+池化
        x = self.conv3(x)
        x = self.bn3(x)  # BN
        x = torch.relu(x)
        x = self.pool3(x)

        # 将特征图做成以为向量的形式：相当于特征向量
        x = x.reshape(x.size(0), -1)

        # 全连接层
        x = torch.relu(self.linear1(x))
        x = self.dropout(x)
        x = torch.relu(self.linear2(x))
        x = self.dropout(x)

        x = torch.relu(self.linear3(x))

        # 返回输出结果
        return self.out(x)


def invert_colors(image_path, output_path=None):
    """
    对输入图片进行颜色反转处理

    参数:
        image_path: 输入图片路径
        output_path: 可选参数，指定反转后图片保存路径。若不提供则只返回处理后的tensor

    返回:
        如果output_path为None，返回处理后的tensor
        否则保存图片并返回None
    """
    # 打开图片并转换为RGB
    image = Image.open(image_path).convert('RGB')

    # 使用PIL直接反转颜色
    inverted_pil = Image.eval(image, lambda x: 255 - x)
    # # 小于126的变为0，大于126的变为255
    # inverted_pil = Image.eval(image, lambda x: 0 if 255 - x < 126 else 255)

    if output_path:
        # 保存反转后的图片
        inverted_pil.save(output_path)
        return None
    else:
        # 转换为tensor返回
        transform = transforms.ToTensor()
        return transform(inverted_pil)




def pro_image(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image)

    return image

def predict_image(image_path, model_path=PATH_best):
    # 对图片进行颜色反转
    invert_colors(image_path, output_path=image_path)

    image = pro_image(image_path).unsqueeze(0).to(device)

    model = ImageClassification().to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()

    # 预测
    with torch.no_grad():
        output = model(image)
        probabilities = torch.softmax(output, dim=1)  # 转换为概率
        pred_class = torch.argmax(probabilities).item()  # 获取预测类别
        confidence = probabilities[0][pred_class].item()  # 获取置信度

    # 0-9数字分类
    class_names = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

    # 返回类别和置信度
    return class_names[pred_class], confidence

# print(predict_image('test_img/5_invert.png'))